// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/exchange_client.h>

#include <pollux/common/base/counters.h>
#include <pollux/common/base/stats_reporter.h>

namespace kumo::pollux::exec {
    void ExchangeClient::addRemoteTaskId(const std::string &remoteTaskId) {
        std::vector<RequestSpec> requestSpecs;
        std::shared_ptr<ExchangeSource> toClose; {
            std::lock_guard<std::mutex> l(queue_->mutex());

            bool duplicate = !remoteTaskIds_.insert(remoteTaskId).second;
            if (duplicate) {
                // Do not add sources twice. Presto protocol may add duplicate sources
                // and the task updates have no guarantees of arriving in order.
                return;
            }

            std::shared_ptr<ExchangeSource> source;
            try {
                source =
                        ExchangeSource::create(remoteTaskId, destination_, queue_, pool_);
            } catch (const PolluxException &) {
                throw;
            } catch (const std::exception &e) {
                // 'remoteTaskId' can be very long. Truncate to 128 characters.
                POLLUX_FAIL(
                    "Failed to create ExchangeSource: {}. Task ID: {}.",
                    e.what(),
                    remoteTaskId.substr(0, 128));
            }

            if (closed_) {
                toClose = std::move(source);
            } else {
                sources_.push_back(source);
                queue_->addSourceLocked();
                emptySources_.push(source);
                requestSpecs = pickSourcesToRequestLocked();
            }
        }

        // Outside of lock.
        if (toClose) {
            toClose->close();
        } else {
            request(std::move(requestSpecs));
        }
    }

    void ExchangeClient::noMoreRemoteTasks() {
        queue_->noMoreSources();
    }

    void ExchangeClient::close() {
        std::vector<std::shared_ptr<ExchangeSource> > sources;
        std::queue<ProducingSource> producingSources;
        std::queue<std::shared_ptr<ExchangeSource> > emptySources; {
            std::lock_guard<std::mutex> l(queue_->mutex());
            if (closed_) {
                return;
            }
            closed_ = true;
            sources = std::move(sources_);
            producingSources = std::move(producingSources_);
            emptySources = std::move(emptySources_);
        }

        // Outside of mutex.
        for (auto &source: sources) {
            source->close();
        }
        queue_->close();
    }

    melon::F14FastMap<std::string, RuntimeMetric> ExchangeClient::stats() const {
        melon::F14FastMap<std::string, RuntimeMetric> stats;
        std::lock_guard<std::mutex> l(queue_->mutex());

        for (const auto &source: sources_) {
            if (source->supportsMetrics()) {
                for (const auto &[name, value]: source->metrics()) {
                    if (UNLIKELY(stats.count(name) == 0)) {
                        stats.insert(std::pair(name, RuntimeMetric(value.unit)));
                    }
                    stats[name].merge(value);
                }
            } else {
                for (const auto &[name, value]: source->stats()) {
                    stats[name].addValue(value);
                }
            }
        }

        stats["peakBytes"] =
                RuntimeMetric(queue_->peakBytes(), RuntimeCounter::Unit::kBytes);
        stats["numReceivedPages"] = RuntimeMetric(queue_->receivedPages());
        stats["averageReceivedPageBytes"] = RuntimeMetric(
            queue_->averageReceivedPageBytes(), RuntimeCounter::Unit::kBytes);

        return stats;
    }

    std::vector<std::unique_ptr<SerializedPage> > ExchangeClient::next(
        int consumerId,
        uint32_t maxBytes,
        bool *atEnd,
        ContinueFuture *future) {
        std::vector<RequestSpec> requestSpecs;
        std::vector<std::unique_ptr<SerializedPage> > pages;
        ContinuePromise stalePromise = ContinuePromise::makeEmpty(); {
            std::lock_guard<std::mutex> l(queue_->mutex());
            if (closed_) {
                *atEnd = true;
                return pages;
            }

            *atEnd = false;
            pages = queue_->dequeueLocked(
                consumerId, maxBytes, atEnd, future, &stalePromise);
            if (*atEnd) {
                return pages;
            }

            if (!pages.empty() && queue_->totalBytes() > maxQueuedBytes_) {
                return pages;
            }

            requestSpecs = pickSourcesToRequestLocked();
        }

        // Outside of lock
        if (stalePromise.valid()) {
            stalePromise.setValue();
        }
        request(std::move(requestSpecs));
        return pages;
    }

    void ExchangeClient::request(std::vector<RequestSpec> &&requestSpecs) {
        auto self = shared_from_this();
        for (auto &spec: requestSpecs) {
            auto future = melon::SemiFuture<ExchangeSource::Response>::makeEmpty();
            if (spec.maxBytes == 0) {
                future = spec.source->requestDataSizes(kRequestDataSizesMaxWaitSec_);
            } else {
                future = spec.source->request(spec.maxBytes, kRequestDataMaxWait);
            }
            POLLUX_CHECK(future.valid());
            std::move(future)
                    .via(executor_)
                    .thenValue(
                        [self, spec = std::move(spec), sendTimeMs = getCurrentTimeMs()](
                    ExchangeSource::Response &&response) {
                            const auto requestTimeMs = getCurrentTimeMs() - sendTimeMs;
                            if (spec.maxBytes == 0) {
                                RECORD_HISTOGRAM_METRIC_VALUE(
                                    kMetricExchangeDataSizeTimeMs, requestTimeMs);
                                RECORD_METRIC_VALUE(kMetricExchangeDataSizeCount);
                            } else {
                                RECORD_HISTOGRAM_METRIC_VALUE(
                                    kMetricExchangeDataTimeMs, requestTimeMs);
                                RECORD_METRIC_VALUE(kMetricExchangeDataBytes, response.bytes);
                                RECORD_HISTOGRAM_METRIC_VALUE(
                                    kMetricExchangeDataSize, response.bytes);
                                RECORD_METRIC_VALUE(kMetricExchangeDataCount);
                            }

                            bool pauseCurrentSource = false;
                            std::vector<RequestSpec> requestSpecs;
                            std::shared_ptr<ExchangeSource> currentSource = spec.source; {
                                std::lock_guard<std::mutex> l(self->queue_->mutex());
                                if (self->closed_) {
                                    return;
                                }
                                if (!response.atEnd) {
                                    if (!response.remainingBytes.empty()) {
                                        for (auto bytes: response.remainingBytes) {
                                            POLLUX_CHECK_GT(bytes, 0);
                                        }
                                        self->producingSources_.push(
                                            {
                                                std::move(spec.source),
                                                std::move(response.remainingBytes)
                                            });
                                    } else {
                                        self->emptySources_.push(std::move(spec.source));
                                    }
                                }
                                self->totalPendingBytes_ -= spec.maxBytes;
                                requestSpecs = self->pickSourcesToRequestLocked();
                                pauseCurrentSource =
                                        std::find_if(
                                            requestSpecs.begin(),
                                            requestSpecs.end(),
                                            [&currentSource](const RequestSpec &spec) -> bool {
                                                return spec.source.get() == currentSource.get();
                                            }) == requestSpecs.end();
                            }
                            if (pauseCurrentSource) {
                                currentSource->pause();
                            }
                            self->request(std::move(requestSpecs));
                        })
                    .thenError(
                        melon::tag_t<std::exception>{}, [self](const std::exception &e) {
                            self->queue_->setError(e.what());
                        });
        }
    }

    std::vector<ExchangeClient::RequestSpec>
    ExchangeClient::pickSourcesToRequestLocked() {
        if (closed_) {
            return {};
        }
        std::vector<RequestSpec> requestSpecs;
        while (!emptySources_.empty()) {
            auto &source = emptySources_.front();
            POLLUX_CHECK(source->shouldRequestLocked());
            requestSpecs.push_back({std::move(source), 0});
            emptySources_.pop();
        }
        int64_t availableSpace =
                maxQueuedBytes_ - queue_->totalBytes() - totalPendingBytes_;
        while (availableSpace > 0 && !producingSources_.empty()) {
            auto &source = producingSources_.front().source;
            int64_t requestBytes = 0;
            for (auto bytes: producingSources_.front().remainingBytes) {
                availableSpace -= bytes;
                if (availableSpace < 0) {
                    break;
                }
                requestBytes += bytes;
            }
            if (requestBytes == 0) {
                POLLUX_CHECK_LT(availableSpace, 0);
                break;
            }
            POLLUX_CHECK(source->shouldRequestLocked());
            requestSpecs.push_back({std::move(source), requestBytes});
            producingSources_.pop();
            totalPendingBytes_ += requestBytes;
        }

        if ((queue_->totalBytes() + totalPendingBytes_ < minOutputBatchBytes_) &&
            !producingSources_.empty()) {
            // Two cases which we request an out-of-band data transfer:
            // 1. We have full capacity but still cannot initiate one single data
            //    transfer. Let the transfer happen in this case to avoid getting stuck.
            //
            // 2. We have some data in the queue that is not big enough for
            //    consumers and it is big enough to not allow ExchangeClient to
            //    initiate request for more data. Let transfer happen in this case
            //    to avoid this deadlock situation.
            auto &source = producingSources_.front().source;
            auto requestBytes = producingSources_.front().remainingBytes.at(0);
            KLOG(INFO) << "Requesting large single page " << requestBytes
              << " bytes, exceeding capacity " << maxQueuedBytes_;
            POLLUX_CHECK(source->shouldRequestLocked());
            requestSpecs.push_back({std::move(source), requestBytes});
            producingSources_.pop();
            totalPendingBytes_ += requestBytes;
        }
        return requestSpecs;
    }

    ExchangeClient::~ExchangeClient() {
        close();
    }

    std::string ExchangeClient::toString() const {
        std::stringstream out;
        for (auto &source: sources_) {
            out << source->toString() << std::endl;
        }
        return out.str();
    }

    melon::Dynamic ExchangeClient::toJson() const {
        melon::Dynamic obj = melon::Dynamic::object;
        obj["taskId"] = taskId_;
        obj["closed"] = closed_;
        melon::Dynamic clientsObj = melon::Dynamic::object;
        int index = 0;
        for (auto &source: sources_) {
            clientsObj[std::to_string(index++)] = source->toJson();
        }
        obj["clients"] = clientsObj;
        return obj;
    }
} // namespace kumo::pollux::exec
